
#include "../layout.h"
#include "../../mds/run.h"
#include "../../mds/io/xyz.h"
#include "../../mds/lj.h"

#include "../../ocl/driver.h"
#include "../../ocl/devmap.h"

#include "../work_group_size.h"

#define MD_ERROR prog_exit = -1; goto END
#define MDS_XERROR(CALL) do{ if ( CALL < 0 ) { mds_perror( stderr, #CALL ); MD_ERROR; } } while(0)
#define OCL_XERROR(CALL) do{ if ( CALL < 0 ) { ocl_perror( stderr, #CALL ); MD_ERROR; } } while(0)

static ocl_driver *    OpenCL = NULL;

static
int on_option_np( const char *value, void *args )
{
	size_t *num_particles = args;
	fprintf( stderr, "-- parsing num_particles\n");
	*num_particles = strtol(value, NULL, 10);
	if( *num_particles <= 0 )
	{
		fprintf(stderr, "invalid num_particles\n");
		return -1;
	}
	return 0;
}

static
int on_option_box( const char *value, void *args )
{
	mds_vector *box = args;
	fprintf( stderr, "-- parsing box\n");
	if( mds_box_parse( box, value )  < 0 )
	{
		mds_perror( stderr, "mds_box_parse('%s')", value );
		return -1;
	}
	return 0;
}

static
int on_option_devmap( const char *value, void *args )
{
	ocl_devmap **devmap = args;
	assert( OpenCL );
	fprintf( stderr, "-- parsing '%s'\n", value);
	*devmap = ocl_devmap_create(OpenCL,value);
	if( NULL == *devmap )
	{
		ocl_perror(stderr,"ocl_devmap_create(%s)",value);
		return -1;
	}
	return 0;
}

static int on_option_dt( const char *value, void *args )
{
	Real *dt = args;
	fprintf( stderr, "-- parsing dt\n");
	assert( args != NULL );
	*dt = strtod( value, NULL);
	if( *dt <= 0 )
	{
		fprintf(stderr, "invalid dt\n");
		return -1;
	}
	return 0;
}

static int on_option_tmax(const char *value, void *args)
{
	Real *tmax = args;
	fprintf( stderr, "-- parsing tmax\n");
	assert( args != NULL );
	*tmax = strtod( value, NULL);
	if( *tmax <= 0 )
	{
		fprintf(stderr, "invalid tmax\n");
		return -1;
	}
	return 0;
}

static int on_option_T( const char *value, void *args )
{
	Real *T = args;
	fprintf( stderr, "-- parsing T\n");
	assert( args != NULL );
	*T = strtod( value, NULL);
	if( *T < 0 )
	{
		fprintf(stderr, "invalid T\n");
		return -1;
	}
	return 0;
}

static int on_option_f( const char *filename, void *args )
{
	mds_options *options = args;
	fprintf( stderr, "-- appending options from file '%s'\n", filename );
	if( mds_options_load( options, filename, "f") < 0 )
	{
		mds_perror( stderr, "mds_options_load(%s)", filename );
		return -1;
	}
	return 0;
}


int main( int argc, char *argv[])
{
	int             prog_exit      = 0;
	mds_frame       frame          = MDS_FRAME_INIT;
	omd_layout      layout         = OMD_LAYOUT_INIT;
	size_t          num_particles  = 0;
	mds_options     options        = MDS_OPTIONS_INIT;
	ocl_devmap     *devmap         = NULL;
	ocl_context     context        = OCL_CONTEXT_INIT;
	ocl_code        sources        = OCL_CODE_INIT;
	ocl_program     program        = OCL_PROGRAM_INIT;
	ocl_kernel      k_predict      = OCL_KERNEL_INIT;
	ocl_kernel      k_correct      = OCL_KERNEL_INIT;
	ocl_kernel      k_dynamics     = OCL_KERNEL_INIT;
	ocl_kernel      k_sum_VW       = OCL_KERNEL_INIT;
	mds_uvector     cr;
	mds_vector_set  position = NULL;
	mds_vector_set  velocity = NULL;
	ocl_queue       Q        = OCL_QUEUE_INIT;
	struct          { Real dt, half_dt; } step_data;
	Real            dt       = 0;
	mds_lj_data     lj;
	Real            epsilon  = 0.995708; /*-- for argon */
	Real            sigma    = 0.3405;   /*-- for argon */
	Real            t_max    = 0;
	Real            t_run    = 0, t_old = t_run;
	Real            temperature = 0;
	mds_wtime       chrono;
	ocl_buffer      dataVW = OCL_BUFFER_INIT;
	ocl_buffer      VW     = OCL_BUFFER_INIT;
	register size_t i;
	
	
	assert( sizeof( step_data ) == sizeof(Real2) );
	assert( sizeof( lj     )    == sizeof(Real8) );
	
	memset( &step_data, 0, sizeof(step_data) );
	mds_run_initialize();
	mds_lj_init( &lj, epsilon, sigma);
	mds_wtime_init( &chrono );
	
	/* ocl_logfile = stderr;
	 */
	
	fprintf( stderr, "-- Loading OpenCL...\n");fflush(stderr);
	OpenCL = ocl_driver_load();
	if( !OpenCL )
	{
		ocl_perror( stderr, "ocl_driver_load()");
		MD_ERROR;
	}
	
	fprintf( stderr, "Ok!\n");fflush(stderr);
	
	MDS_OPTIONS_DECL()
	{
		MDS_OPTION_NORMAL(f,      options,      "read options from file"),
		MDS_OPTION_FORCED(np,     num_particles,"particles in simulation"),
		MDS_OPTION_FORCED(box,    frame.box,    "simulation box"),
		MDS_OPTION_FORCED(devmap, devmap,       "OpenCL devices map"),
		MDS_OPTION_FORCED(dt,     dt,           "Simulation Time Step"),
		MDS_OPTION_FORCED(tmax,   t_max,        "Max Simulation Time"),
		MDS_OPTION_NORMAL(T,temperature,        "temperature (default=0)")
	}
	MDS_OPTIONS_DONE();
	
	frame.pbc_for_f = mds_pbc_xyz;
	frame.pbc_for_r = mds_pbc_xyz;
	
	if( argc <= 1 )
	{
		MDS_OPTIONS_HELP(stderr);
		return 1;
	}
	
	fprintf( stderr, "-- parsing options...\n" );
	MDS_XERROR(  mds_options_parse(&options, argc, argv) );
	fprintf( stderr, "-- dispatching options...\n" );
	MDS_OPTIONS_DISPATCH(options,MD_ERROR);	
	fprintf( stderr, "-- done!\n");
	fflush(stderr);
	assert( devmap );
	
	/***************************************************************************
	 *
	 * adjust particles to crystal
	 *
	 **************************************************************************/
	fprintf( stderr, "-- counting particles on fcc...\n"); fflush(stderr);
	num_particles = mds_crystal_close_packed_count( &cr, num_particles, & frame.box);
	fprintf( stderr, "-- particles: %u\n", (unsigned) num_particles ); fflush(stderr);
	
	/** create context **/
	fprintf( stderr, "-- creating context\n"); fflush(stderr);
	OCL_XERROR( ocl_context_map( &context, devmap) );
	
	/** load sources **/
	fprintf( stderr, "-- reading code\n"); fflush(stderr);
	OCL_XERROR( ocl_code_read( &sources, "kernels.cl" ) );
	
	
	/** create program **/
	fprintf( stderr, "-- create program\n"); fflush(stderr);
	OCL_XERROR( ocl_program_load_sources( &program, &context, &sources) );
	
	
	/** compile program **/
	fprintf( stderr, "-- compile program\n"); fflush(stderr);
	if( ocl_program_compile( &program, ocl_options) < 0 )
	{
		fprintf( stderr, "-- FAILURE\n"); fflush(stderr);
		ocl_perror( stderr, "ocl_compile_program(%s)", ocl_options);
		ocl_program_display_build_logs( &program, stderr);
		MD_ERROR;
	}
	else 
	{
		fprintf( stderr, "-- SUCCESS\n"); fflush(stderr);
		ocl_program_display_build_logs( &program, stderr);
	}
	
	/** load kernels **/
	fprintf( stderr, "-- loading kernels\n"); fflush(stderr);
	OCL_XERROR( ocl_kernel_load( &k_predict,  &program, "predict"  ) );
	OCL_XERROR( ocl_kernel_load( &k_correct,  &program, "correct"  ) );
	OCL_XERROR( ocl_kernel_load( &k_dynamics, &program, "dynamics" ) );
	OCL_XERROR( ocl_kernel_load( &k_sum_VW,   &program, "sum_VW"   ) );
	
	/** create frame, need only 2 vector_sets **/
	fprintf( stderr, "-- creating frame\n"); fflush(stderr);
	MDS_XERROR( mds_frame_load( &frame, num_particles, 2));
	
	/** create layout **/
	fprintf( stderr, "-- creating layout\n"); fflush(stderr);
	OCL_XERROR( omd_layout_load( &layout, &context, num_particles, 3));
	
	
	step_data.dt      = dt;
	step_data.half_dt = dt/2;
	
	
	
	/** setting kernels arguments **/
	OCL_KERNEL_SET_ARGS_INIT(&k_predict,arg_index,arg_result);
	OCL_KERNEL_SET_ARG_BUFFER(& layout.buffers[MDS_POSITION]     );
	OCL_KERNEL_SET_ARG_BUFFER(& layout.buffers[MDS_VELOCITY]     );
	OCL_KERNEL_SET_ARG_BUFFER(& layout.buffers[MDS_ACCELERATION] );
	OCL_KERNEL_SET_ARG(Real2,step_data);
	OCL_KERNEL_SET_ARG(Real4,frame.box);
	OCL_XERROR( arg_result );
	OCL_KERNEL_SET_ARGS_DONE();
	
	/** setting kernels arguments **/
	OCL_KERNEL_SET_ARGS_INIT(&k_correct,arg_index,arg_result);
	OCL_KERNEL_SET_ARG_BUFFER(& layout.buffers[MDS_POSITION]     );
	OCL_KERNEL_SET_ARG_BUFFER(& layout.buffers[MDS_VELOCITY]     );
	OCL_KERNEL_SET_ARG_BUFFER(& layout.buffers[MDS_ACCELERATION] );
	OCL_KERNEL_SET_ARG(Real2,step_data);
	OCL_KERNEL_SET_ARG(Real4,frame.box);
	OCL_XERROR( arg_result );
	OCL_KERNEL_SET_ARGS_DONE();
	
	/** setting kernels arguments **/
	OCL_KERNEL_SET_ARGS_INIT(&k_dynamics,arg_index,arg_result);
	OCL_KERNEL_SET_ARG_BUFFER(& layout.buffers[MDS_POSITION]     );
	OCL_KERNEL_SET_ARG_BUFFER(& layout.buffers[MDS_VELOCITY]     );
	OCL_KERNEL_SET_ARG_BUFFER(& layout.buffers[MDS_ACCELERATION] );
	OCL_KERNEL_SET_ARG(Real8,lj);
	OCL_KERNEL_SET_ARG(Real4,frame.box);
	OCL_XERROR( arg_result );
	OCL_KERNEL_SET_ARGS_DONE();
	
	position = frame.sets[MDS_POSITION];
	velocity = frame.sets[MDS_VELOCITY];
	
	for( i=0;i<num_particles;++i )
	{
		const Real mass           = 10 + mds_alea();
		velocity[i].info.mass     = mass;
		position[i].info.inv_mass = MDS_F(1.0)/mass;
		
	}
	
	mds_crystal_build_fcc( position, num_particles, &cr, &frame.box, 0);
	mds_frame_initialize_velocities( &frame, temperature );
	
	(void) mds_xyz_save( "md.xyz", position, num_particles, "frame", 0);
	fprintf( stderr, "K0=%g\n", frame.K );
	
	{
		const cl_device_id device_id = context.DEVICES[0];
		const ocl_device  *device    = ocl_driver_device( device_id );
		register size_t    iter = 0;
		const size_t       global_work_size[] = { num_particles };
		size_t             local_work_size[]  = { num_particles };
		size_t             num_work_groups    = 0;
		double             total_ell = 0;
		fprintf( stderr, "-- on '%s'\n", device->NAME );
		
		/** find the best work_group_size for dynamics */
		OCL_XERROR( omd_work_group_size_optimize( &k_dynamics, device, global_work_size[0], local_work_size) );
		fprintf( stderr, "-- global_work_size= %u, local_work_size= %u\n", (unsigned) global_work_size[0], (unsigned) local_work_size[0] );
		
		/** count the #work_groups **/
		num_work_groups = global_work_size[0] / local_work_size[0];
		fprintf( stderr, "-- creating dataIO for #work_groups= %u \n", (unsigned)num_work_groups); fflush( stderr );
		OCL_XERROR( ocl_buffer_load( &dataVW, &context, CL_MEM_READ_WRITE, num_work_groups * sizeof(Real2), NULL) );
		OCL_XERROR( ocl_buffer_load( &VW,     &context, CL_MEM_READ_WRITE, sizeof(Real2), NULL) );
		
		OCL_XERROR( ocl_kernel_set_arg_buffer( &k_sum_VW, 0, &dataVW) );
		OCL_XERROR( ocl_kernel_set_arg_buffer( &k_sum_VW, 1, &VW ) );
		{
			const cl_int nwg = num_work_groups;
			OCL_XERROR( ocl_kernel_set_arg(&k_sum_VW, 2, sizeof(nwg), &nwg) );
		}
		
		/** assign local memory */
		OCL_XERROR( ocl_kernel_set_arg_buffer( &k_dynamics, k_dynamics.NUM_ARGS-3, &dataVW) );
		fprintf(stderr, "-- assign local memory\n");
		OCL_XERROR( ocl_kernel_set_arg_local(  &k_dynamics, k_dynamics.NUM_ARGS-2, local_work_size[0] * sizeof(Real4)     ) );
		OCL_XERROR( ocl_kernel_set_arg_local(  &k_dynamics, k_dynamics.NUM_ARGS-1, local_work_size[0] * (2*sizeof(Real4)) ) );
		
		/** create the queue **/
		
		OCL_XERROR( ocl_queue_load( &Q, &context, device_id, 0) );
		/** send initial data to device **/
		fprintf( stderr, "-- send initial condition to device\n");
		OCL_XERROR( omd_layout_write( &layout, &Q, &frame, MDS_POSITION, CL_FALSE) );
		OCL_XERROR( omd_layout_write( &layout, &Q, &frame, MDS_VELOCITY, CL_FALSE) );
		
		/** first dynamics **/
		OCL_XERROR( ocl_enqueue_ndrange_kernel( &Q, &k_dynamics, 1, NULL, global_work_size, local_work_size, OCL_NO_EVENTS) );
		
		/** fetch V and W **/
		OCL_XERROR( ocl_enqueue_task( &Q, &k_sum_VW, OCL_NO_EVENTS) );
		OCL_XERROR( ocl_enqueue_read_buffer(&Q, &VW, CL_TRUE, 0, sizeof(Real2), & frame.V, OCL_NO_EVENTS) );
		
		fprintf( stderr, "Initial | V=%8.4f | K= %8.4f\n", frame.V, frame.K);
		{
			FILE *fp = fopen( "energy.txt", "wt" );
			if( fp )
			{
				fprintf( fp, "%.6e %.6e %.6e %.6e\n", 0.0, frame.V, frame.K, frame.V + frame.K );
				fclose( fp );
			}
		}
		
		mds_wtime_start( &chrono );
		while( t_run < t_max )
		{
			t_run = ++iter * dt;
			OCL_XERROR( ocl_enqueue_ndrange_kernel(&Q, &k_predict,  1, NULL, global_work_size, NULL,            OCL_NO_EVENTS) );
			OCL_XERROR( ocl_enqueue_ndrange_kernel(&Q, &k_dynamics, 1, NULL, global_work_size, local_work_size, OCL_NO_EVENTS) );
			OCL_XERROR( ocl_enqueue_ndrange_kernel(&Q, &k_correct,  1, NULL, global_work_size, NULL,            OCL_NO_EVENTS) );
			
			
			if( t_run - t_old >= 0.05 )
			{
				OCL_XERROR( omd_layout_read( &layout, &Q, &frame, MDS_POSITION,     CL_FALSE) );
				OCL_XERROR( omd_layout_read( &layout, &Q, &frame, MDS_VELOCITY,     CL_FALSE) );
				/** fetch V and W **/
				/*
				 * OCL_XERROR( ocl_enqueue_task( &Q, &k_sum_VW, OCL_NO_EVENTS) );
				OCL_XERROR( ocl_enqueue_read_buffer(&Q, &VW, CL_FALSE, 0, sizeof(Real2), & frame.V, OCL_NO_EVENTS) );
				*/
				
				OCL_XERROR( ocl_finish( &Q ) );
				frame.K = mds_velocity_kinetic_energy( velocity, num_particles);
				
				total_ell += mds_wtime_query( &chrono );
				fprintf( stderr, "t=  %8.4f | V=%.6e | K= %.6e | E = %.6e <%8.4f> fps        \r", t_run, frame.V, frame.K, frame.V + frame.K, ((double)iter)/total_ell); fflush( stderr );
				(void) mds_xyz_save( "md.xyz", position, num_particles, "frame", 1);
				{
					FILE *fp = fopen( "energy.txt", "at" );
					if( fp )
					{
						fprintf( fp, "%.6e %.6e %.6e %.6e\n", t_run, frame.V, frame.K, frame.V + frame.K );
						fclose( fp );
					}
				}
				t_old = t_run;
				mds_wtime_start( &chrono );
			}
		}
		fprintf( stderr, "\n\n");
		
		
		OCL_XERROR( ocl_finish( &Q ) );
	}
	goto END;
END:
	MDS_MEMORY();
	ocl_buffer_quit( &VW     );
	ocl_buffer_quit( &dataVW );
	ocl_queue_quit( &Q );
	ocl_kernel_quit( &k_sum_VW   );
	ocl_kernel_quit( &k_dynamics );
	ocl_kernel_quit( &k_correct );
	ocl_kernel_quit( &k_predict );
	ocl_program_quit( &program  );
	ocl_code_quit( &sources );
	ocl_context_quit( &context );
	ocl_devmap_delete( devmap );
	mds_options_quit( &options );
	omd_layout_quit( &layout );
	ocl_context_quit( &context );
	mds_frame_quit(  &frame  );
	ocl_driver_quit();
	MDS_MEMORY();
	return prog_exit;
}
